
import argparse
import logging
import time
import sys
from model_manage import ModelManager
from parse import Parser
from config import config_singleton, config_model_singleton
from udp_handle import Udp_Multicast_Receiver
from alloter import Alloter
from threading import Thread
from utils import sys_logger
import keyboard
from collector import Collector
from tcp_handle import TcpClient


if __name__ == '__main__':

    """

    """
    # 设置输入参数
    parser = argparse.ArgumentParser(description="Anomaly Detection System for Rocket Test Data.")
    parser.add_argument('-t', "--train", default=None, required=False)

    need_stop_main = False

    # 创建udp组播接收,
    udp_ip = config_singleton.udp_setting["ip"]
    udp_port = config_singleton.udp_setting["port"]
    udp_maxsize = config_singleton.udp_setting["queue_maxsize"]
    udp_recv = Udp_Multicast_Receiver(udp_ip, udp_port, udp_maxsize)
    sys_logger.info("UDP Multicast Receiver created.")

    # 创建TCP或者UDP客户端,用于数据发送
    tcp_ip = config_singleton.tcp_setting["ip"]
    tcp_port = config_singleton.tcp_setting["port"]
    tcp_maxsize = config_singleton.tcp_setting["queue_maxsize"]
    as_udp = config_singleton.tcp_setting["as_udp"]
    
    # 根据配置，使用udp发送还是用tcp发送
    if as_udp:
        send_client = Udp_Multicast_Receiver(tcp_ip, tcp_port, tcp_maxsize, True)
        sys_logger.info("UDP Multicast Send Client created.")
    else:
        send_client = TcpClient(tcp_ip, tcp_port, tcp_maxsize)
        send_client.connect()

        if not send_client.connected:
            sys_logger.error("Failed to connect to TCP server. Exiting...")
            sys.exit()

    # 创建parser
    parser_queue_size = config_singleton.parse_configs["max_msgbean_size"]
    drop_bytes = config_singleton.parse_configs["drop_bytes_everytime"]
    parser = Parser(udp_recv=udp_recv, 
                    max_msgbean_size=parser_queue_size, 
                    drop_bytes=drop_bytes)
    sys_logger.info("Parser Created.")

    # 创建模型管理器
    model_configs = config_model_singleton.model_configs
    model_manager = ModelManager(model_configs)
    model_manager.create_models()   # 创建模型
    sys_logger.info("Model Manager Created.") 

    # 创建分配器
    alloter = Alloter(model_manager, parser)
    sys_logger.info("Alloter Created.")

    # 创建数据收集器
    collector = Collector(writetoboard=config_singleton.write_to_tensorboard,
                          base_path=config_singleton.tensorboard_log_root,
                          manager=model_manager,
                          handler=send_client)

    sys_logger.info("Collector Created.")
    sys_logger.info("====================================================")

    # 创建线程
    alloter_thread = Thread(target=alloter.allot_run)
    alloter_thread.start()
    sys_logger.info("Alloter Thread Started.")

    parse_thread = Thread(target=parser.do_parse)
    parse_thread.start()
    sys_logger.info("Parser Thread Started.") 

    # 启动所有模型预测线程
    model_manager.run_all_models()
    sys_logger.info("All Model Threads Started.")

    # 启动tcp发送线程
    send_thread = Thread(target=send_client.do_send)
    send_thread.start()
    sys_logger.info("Send Thread Started.")


    # 启动数据收集线程
    collector_thread = Thread(target=collector.do_collect)
    collector_thread.start()
    sys_logger.info("Collector Thread Started.")
    

    # 启动UDP Multicast Receiver Thread Started
    udp_recv_thread = Thread(target=udp_recv.do_receive)
    udp_recv_thread.start()
    sys_logger.info("UDP Multicast Receiver Thread Started.")
    sys_logger.info("====================================================")

    # 设置监听快捷键的响应函数
    def stop_main():

        udp_recv.stop()
        time.sleep(1)

        parser.stop()
        time.sleep(1)

        alloter.stop_run()
        time.sleep(1)

        model_manager.stop_all_models()
        time.sleep(1)

        collector.stop()
        time.sleep(1)

        send_client.stop()
        time.sleep(1)

        sys_logger.info("System Shutdown Finished.")
        sys_logger.setLevel(logging.NOTSET)  # 停止logging
        print("****PROGRAM END****")


    keyboard.add_hotkey('ctrl+shift+q', stop_main)
    print("Press ***Ctrl+Shift+Q*** to stop the system.")


